/*
 *         (C) COPYRIGHT NBI Limited
 *              ALL RIGHT RESERVED
 *
 * File        : fastflow/graph/model.cc
 * Authors     : dzhang
 * Create Time : 2021/09/04 17:25:08
 * Description :
 *
 */

#include "fastflow/graph/model.h"
#include "fastflow/core/logging.h"
#include "fastflow/graph/schema_registry.h"

using namespace ONNX_NAMESPACE;

namespace fastflow {


Model::Model(const std::string& graph_name,
             bool is_onnx_domain_only,
             const ModelMetaData& model_metadata,
             const IFastFlowOpSchemaRegistryList local_registries,
             const std::unordered_map<std::string, int>& domain_to_version,
             const std::vector<ONNX_NAMESPACE::FunctionProto>& model_functions) {
  model_proto_ = std::make_unique<ModelProto>();
  model_proto_->set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION);
  model_proto_->mutable_graph()->set_name(graph_name);
  model_metadata_ = model_metadata;
  for (auto& metadata : model_metadata_) {
    const gsl::not_null<StringStringEntryProto*> prop{model_proto_->add_metadata_props()};
    prop->set_key(metadata.first);
    prop->set_value(metadata.second);
  }

  auto schema_registry = std::make_shared<SchemaRegistryManager>();
  for (auto schema_collection : local_registries) {
    schema_registry->RegisterRegistry(schema_collection);
  }

  auto* p_domain_to_version = &domain_to_version;
  std::unordered_map<std::string, int> domain_to_version_static;
  if (p_domain_to_version->empty()) {
    domain_to_version_static = schema_registry->GetLatestOpsetVersions(is_onnx_domain_only);
    p_domain_to_version = &domain_to_version_static;
  }

  for (auto domain : *p_domain_to_version) {
    const gsl::not_null<OperatorSetIdProto*> opset_id_proto{model_proto_->add_opset_import()};
    opset_id_proto->set_domain(domain.first);
    opset_id_proto->set_version(domain.second);
  }

  std::unordered_map<std::string, const ONNX_NAMESPACE::FunctionProto*> model_functions_map;
//   for (auto& func : model_functions) {
//     auto func_ptr = model_proto_->add_functions();
//     func_ptr->CopyFrom(func);
//     model_functions_map[func_ptr->name()] = func_ptr;
//   }

  // need to call private ctor so can't use make_shared
  graph_.reset(new Graph(model_proto_->mutable_graph(), *p_domain_to_version, IrVersion(), schema_registry, model_functions_map));
}

Model::Model(const ModelProto& model_proto, const IFastFlowOpSchemaRegistryList* local_registries)
    : Model(std::make_unique<ModelProto>(model_proto), local_registries) {
}

Model::Model(std::unique_ptr<ModelProto> model_proto, const IFastFlowOpSchemaRegistryList* local_registries) {
  if (!model_proto) {
    throw std::invalid_argument("ModelProto was null.");
  }

  if (!model_proto->has_graph()) {
    throw std::invalid_argument("ModelProto does not have a graph.");
  }

  if (model_proto->opset_import_size() == 0) {
    throw std::invalid_argument(
        "Missing opset in the model. All ModelProtos MUST have at least one entry that"
        " specifies which version of the ONNX OperatorSet is being imported.");
  }

  model_proto_.reset(std::move(model_proto.release()));
  for (auto& prop : model_proto_->metadata_props()) {
    model_metadata_[prop.key()] = prop.value();
  }

  auto schema_registry = std::make_shared<SchemaRegistryManager>();
  if (local_registries != nullptr) {
    for (auto schema_collection : *local_registries) {
      schema_registry->RegisterRegistry(schema_collection);
    }
  }

  std::unordered_map<std::string, int> domain_to_version;
  for (auto& opSet : model_proto_->opset_import()) {
    domain_to_version[opSet.domain()] = gsl::narrow_cast<int>(opSet.version());
  }

  auto domain_map = schema_registry->GetLatestOpsetVersions(false);
  for (auto domain : domain_map) {
    if (domain_to_version.find(domain.first) == domain_to_version.end()) {
      domain_to_version[domain.first] = domain.second;
      const gsl::not_null<OperatorSetIdProto*> opset_id_proto{model_proto_->add_opset_import()};
      opset_id_proto->set_domain(domain.first);
      opset_id_proto->set_version(domain.second);
    }
  }

  std::unordered_map<std::string, const ONNX_NAMESPACE::FunctionProto*> model_functions_map;
  // for (auto& func : model_proto_->functions()) {
  //   model_functions_map[func.name()] = &func;
  // }

  // create instance. need to call private ctor so can't use make_unique
  graph_.reset(new Graph(model_proto_->mutable_graph(), domain_to_version, IrVersion(), schema_registry, model_functions_map));
  graph_->Resolve();
}

Version Model::IrVersion() const {
  if (model_proto_->has_ir_version()) {
    return model_proto_->ir_version();
  }
  return kNoVersion;
}

Graph& Model::MainGraph() noexcept {
  return *graph_;
}

const Graph& Model::MainGraph() const noexcept {
  return *graph_;
}

}  // namespace fastflow
